import base64
import json
import logging
import os
import time
import uuid

from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework import status
from rest_framework import viewsets

from comfyui import serializers, models, comfyui_cache
from comfyui import work_tool
from common import translate_lang

logger = logging.getLogger("default")


def create_client_id(user_obj=None):
    # c0f0573850714ec8a9c8634ac9b29b6a
    def get_uuid_hex(exclude_list):
        _uuid = uuid.uuid4().hex
        if _uuid in exclude_list:
            _uuid = get_uuid_hex(exclude_list)
        else:
            return _uuid

    if user_obj:
        client_ids = models.ComfyuiClient.objects.all().values_list('client_id', flat=True)
        _client_id = get_uuid_hex(client_ids)
        _new = models.ComfyuiClient.objects.create(client_id=_client_id, user=user_obj)
    else:
        _client_id = "c0f0573850714ec8a9c8634ac9b29b6a"
    return _client_id


# Create your views here.
class ComfyuiAPIView(viewsets.GenericViewSet):

    @action(methods=['post'], detail=False, permission_classes=())
    def generate_image(self, request):
        """文生图"""
        serializer = serializers.TextToImgSerialization(data=request.data)
        if not serializer.is_valid():
            return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
        validated_data = serializer.validated_data
        # print(validated_data)
        # return Response([], status=status.HTTP_200_OK)

        logger.info("validated_data = {}".format(validated_data))
        prompt_text = translate_lang.translate_zn_to_en(validated_data['prompt'])  # 中->英翻译
        negative_prompt_text = translate_lang.translate_zn_to_en(validated_data['negative_prompt']) # 中->英翻译
        workflow_id = validated_data['workflow_id']

        # 先获取对应用户的client_id (数据库)
        client_obj = models.ComfyuiClient.objects.filter(user=request.user)
        client_id = client_obj.first().client_id if client_obj else create_client_id()

        # 一. 获取工作流模板
        workflow_obj = None
        if workflow_id:
            workflow_obj = models.WorkflowInfo.objects.filter(id=workflow_id)
        request_prompt_data = json.loads(workflow_obj.first().workflow) if workflow_obj else work_tool.get_default_workflow()
        # 二. 替换文本和client_id
        # 1.替换extra_data中的工作流workflow的文本节点内容

        # 先获取到K采样器
        ks_positive_link = list()
        ks_negative_link = list()
        input_link_dict = dict()
        output_link_dict = dict()
        extra_data = request_prompt_data['extra_data']
        workflow = extra_data['extra_pnginfo']['workflow']
        for wf_key, wf_value in workflow.items():
            if wf_key == 'nodes':
                for wf_node in wf_value:  # wf_value = [{"id": 6, "type": "CLIPTextEncode", ...}, ]

                    for _ in wf_node.get('outputs', []):
                        for _link in _['links']:
                            output_key = '{}-{}'.format(str(wf_node['id']), wf_node['type'])
                            output_link_dict[_link] = output_key

                    for _ in wf_node.get('inputs', []):
                        input_key = '{}-{}'.format(str(wf_node['id']), wf_node['type'])
                        input_link_dict[_['link']] = input_key
                        # 获取采样器的正向参数对象和负面参数对象
                        if 'KSampler' in wf_node['type']:
                            if _['name'] == 'positive' and _['type'] == 'CONDITIONING':
                                ks_positive_link.append(_['link'])
                            if _['name'] == 'negative' and _['type'] == 'CONDITIONING':
                                ks_negative_link.append(_['link'])

        if not ks_positive_link:
            raise

        positive_nodes = list()
        negative_nodes = list()
        already_check_nodes = list()

        def get_txt_nodes(_link_num, _text_nodes):
            previous_link_node = output_link_dict.get(_link_num, None)
            _node_type = previous_link_node.split('-')
            already_check_nodes.append(_node_type[0])
            # 若追溯到是K采样器或者追溯过了的节点，则结束后续追溯。
            if 'KSampler' in _node_type[1] and _link_num in already_check_nodes:
                return list(set(_text_nodes))

            if previous_link_node and any(['CLIPTextEncode' in _node_type[1], 'PrimitiveNode' in _node_type[1]]):
                _text_nodes.append(_node_type[0])
                # 继续往前追溯
                keys = [k for k, v in input_link_dict.items() if v == previous_link_node]
                for k in keys:
                    _text_nodes.extend(get_txt_nodes(k, _text_nodes))

            return list(set(_text_nodes))

        for _positive_link in ks_positive_link:
            positive_nodes.extend(get_txt_nodes(_positive_link, positive_nodes))

        for _negative_link in ks_negative_link:
            negative_nodes.extend(get_txt_nodes(_negative_link, negative_nodes))

        # 2.替换prompt中的文本节点内容
        prompt = request_prompt_data['prompt']
        for _, node_data in prompt.items():
                if _ in positive_nodes:
                    node_data['inputs']['text'] = prompt_text
                if _ in negative_nodes:
                    node_data['inputs']['text'] = negative_prompt_text
        # 3.组合为新的工作流
        result_workflow = {
            "client_id": client_id,
            "extra_data": extra_data,
            "prompt": prompt
        }
        # 三. 请求comfyui服务的prompt接口
        result = work_tool.request_comfyui_prompt(workflow_json=result_workflow)
        result['client_id']=client_id
        # print("=================result========================")
        # print(result)  # {'prompt_id': 'b1a41f4b-4888-4e30-a990-2c29220356d9', 'number': 3, 'node_errors': {}}
        return Response(result, status=status.HTTP_200_OK)

    @action(methods=['post'], detail=False, permission_classes=())
    def get_prompt_image(self, request):
        """获取结果"""
        serializer = serializers.GetPromptSerializer(data=request.data)
        if not serializer.is_valid():
            return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
        validated_data = serializer.validated_data
        logger.info("validated_data = {}".format(validated_data))
        prompt_id = validated_data['prompt_id']
        client_id = validated_data['client_id']

        comfyui_output_path = "G:\ComfyUI_windows_portable\ComfyUI\output"  # fixme: 从数据库里获取comfyui的安装路径和图片存储路径

        time_out = 5*60
        start_time = time.time()
        encoded_images = list()

        while True:
            result_data = work_tool.get_history(prompt_id)
            prompt_id_data = result_data.get(prompt_id)
            if not prompt_id_data:
                queue_data = work_tool.get_queue()
                queue_running = queue_data.get('queue_running', list())  # 运行中
                queue_pending = queue_data.get('queue_pending', list())  # 待定
                check_running = any([True for _running_data in queue_pending if prompt_id in _running_data])
                check_pending = any([True for _running_data in queue_running if prompt_id in _running_data])
                if check_running or check_pending:
                    print("================ RUNNING ==============", prompt_id)
                    time.sleep(3)
                    continue
            output_data = prompt_id_data.get('outputs')
            # print("=========== output_data = ", output_data)
            # output_data = {'9': {'images': [{'filename': 'Result_00013_.png', 'subfolder': '', 'type': 'output'}]}}
            _break_out = list()
            for node_str, node_info in output_data.items():

                if all([True for _data in node_info.get('images', []) if _data['type'] == 'output']):
                    for image_data in node_info.get('images', []):
                        # 图片路径
                        image_path = os.path.join(comfyui_output_path, image_data['filename'])
                        # print("====================image_path =", image_path)
                        # 读取图片并进行 Base64 编码
                        with open(image_path, "rb") as image_file:
                            encoded_string = str(base64.b64encode(image_file.read()), encoding="utf-8")
                            # encoded_string = bytes.decode(base64.b64encode(image_file.read()))
                            mime_type = "image/png"  # 确保这里的 MIME 类型与你的图片类型匹配
                            data_url = f"data:image/png;base64,{encoded_string}"
                            encoded_images.append(data_url)

                    _break_out.append(True)  # 当所有节点的图片都生成完成后，就退出while循环
            if all(_break_out):
                break

            now_time = time.time()
            if (now_time - start_time) > time_out:
                break
            time.sleep(3)
        # 缓存到历史生图
        client_id_cache = comfyui_cache.CACHE_CLIENT_IMAGE.get(client_id, list())
        if client_id_cache:
            client_id_cache.extend(encoded_images)
            # 资源限制：缓存对象最多保留10条数据
            if len(client_id_cache) >= 10:
                comfyui_cache.CACHE_CLIENT_IMAGE[client_id] = client_id_cache[-10:]
            else:
                comfyui_cache.CACHE_CLIENT_IMAGE[client_id] = client_id_cache
        else:
            comfyui_cache.CACHE_CLIENT_IMAGE[client_id] = encoded_images
        return Response(encoded_images, status=status.HTTP_200_OK)

    @action(methods=['post'], detail=False, permission_classes=())
    def history(self, request):
        serializer = serializers.GenerationHistorySerializer(data=request.data)
        if not serializer.is_valid():
            return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)

        validated_data = serializer.validated_data
        logger.info("validated_data = {}".format(validated_data))
        client_id = validated_data['client_id']
        history_images = comfyui_cache.CACHE_CLIENT_IMAGE.get(client_id, list())[::-1]  # [::-1]反转列表元素
        return Response(history_images)

    @action(methods=['post'], detail=False, permission_classes=())
    def workflows(self, request):
        workflow_obj = models.WorkflowInfo.objects.all().values('id', 'description')
